import cv2 as cv
import numpy as np
import os
from python_ai.common.xcommon import *

VER = 'v1.0'
N_HIS = 500
ALPHA = 1e-3
FPS = 25
INTERVAL = 1000 // FPS
BASE_DIR, FILE_NAME = os.path.split(__file__)
path = '../../../../../large_data/video/traffic/traffic_20s.avi'
VIDEO_PATH = os.path.join(BASE_DIR, path)
SAVE_DIR = os.path.join(BASE_DIR, '_save', FILE_NAME, VER)
LOG_DIR = os.path.join(BASE_DIR, '_log', FILE_NAME, VER)
SAVE_PATH = os.path.join(SAVE_DIR, 'tgt_' + rand_name_on_now() + '.avi')

video = cv.VideoCapture(VIDEO_PATH)
bg = cv.createBackgroundSubtractorMOG2(N_HIS, detectShadows=False)
cnt = 0
print('Training ...')
while True:
    ret, img = video.read()
    if not ret:
        break
    cnt += 1
    if cnt % 25 == 0:
        print(f'Processed {cnt} frames.')
    if cnt > N_HIS:
        break
    bg.apply(img, None, ALPHA)
video.release()
print('Trained.')

video = cv.VideoCapture(VIDEO_PATH)
cnt = 0
kernel = np.ones((3, 3), dtype=np.float64)
fourcc = cv.VideoWriter_fourcc(*list('XVID'))
os.makedirs(os.path.split(SAVE_PATH)[0], exist_ok=True)
tgt = cv.VideoWriter(SAVE_PATH, fourcc, FPS, img.shape[:2][::-1])
print(f'Video writing to {SAVE_PATH}')
while True:
    ret, img = video.read()
    if not ret:
        print('Video over.')
        break
    cnt += 1
    if 0xFF & cv.waitKey(INTERVAL) == 27:
        print('Cancelled by user.')
        break
    mask = bg.apply(img, None, 0.)
    mask = cv.morphologyEx(mask, cv.MORPH_OPEN, kernel)
    mask = cv.morphologyEx(mask, cv.MORPH_CLOSE, kernel)
    mask = cv.dilate(mask, kernel, iterations=2)
    img = cv.bitwise_and(img, img, mask=mask)
    tgt.write(img)
    cv.imshow('traffic', img)
print(f'{cnt} frames.')
video.release()
tgt.release()
print(f'Video written to {SAVE_PATH}')
